// Copyright (c) 2025 Inochi Amaoto

/* SPDX-License-Identifier: GPL-2.0 */

#define pr_fmt(fmt) "vhost-vapp: " fmt

#include <linux/module.h>
#include <linux/scatterlist.h>
#include <linux/vhost.h>
#include <linux/hashtable.h>
#include "vhost-vapp.h"
#include "vhost-vapp-dev.h"
#include "vhost-vapp-dtq.h"

static DEFINE_MUTEX(vhost_vapp_mutex);
static DEFINE_READ_MOSTLY_HASHTABLE(vhost_vapp_hash, 8);

static struct vhost_vapp_dev* __vhost_vapp_get_dev(u64 uuid)
{
  struct vhost_vapp_dev* vapp;
  hash_for_each_possible_rcu(vhost_vapp_hash, vapp, node, uuid) {
    u32 oid = vapp->conf.uuid;

    /* Skip instances that have no uuid yet */
    if (oid == 0)
      continue;

    if (oid == uuid)
      return vapp;
  }

  return NULL;
}

struct vhost_vapp_dev* vhost_vapp_get_dev(u64 uuid)
{
  struct vhost_vapp_dev* vapp;
  mutex_lock(&vhost_vapp_mutex);
  vapp = __vhost_vapp_get_dev(uuid);
  mutex_unlock(&vhost_vapp_mutex);
  return vapp;
}


static void vhost_vapp_flush(struct vhost_vapp_dev* vapp)
{
  size_t i;
  if (vapp->ddtqs) {
    vhost_vapp_dtq_flush(&vapp->cdtq);
    for (i = 0; i < vapp->conf.multiqueues; i++)
      vhost_vapp_dtq_flush(&vapp->ddtqs[i]);
    vhost_work_dev_flush(&vapp->dev);
  }
}

static inline long vhost_vapp_set_features(struct vhost_vapp_dev* vapp, unsigned long arg)
{
  size_t i;
  u64 features;
  if (copy_from_user(&features, (void*)arg, sizeof(features)))
    return -EFAULT;

  if (features & ~VHOST_VAPP_FEATURES)
    return -EOPNOTSUPP;

  mutex_lock(&vapp->dev.mutex);
  if ((features & (1 << VHOST_F_LOG_ALL)) && !vhost_log_access_ok(&vapp->dev)) {
    goto fail;
  }

  if (!vapp->ddtqs)
    goto fail;

  vhost_vapp_dtq_set_features(&vapp->cdtq, features);
  for (i = 0; i < vapp->conf.multiqueues; i++)
    vhost_vapp_dtq_set_features(&vapp->ddtqs[i], features);

  mutex_unlock(&vapp->dev.mutex);
  return 0;

 fail:
  mutex_unlock(&vapp->dev.mutex);
  return -EPERM;
}

static long vhost_vapp_dev_set_uuid(struct vhost_vapp_dev* vapp, unsigned long arg)
{
  u64 uuid;
  struct vhost_vapp_dev* other;
  if (copy_from_user(&uuid, (void*)arg, sizeof(uuid)))
    return -EFAULT;

  if (uuid == 0)
    return -EPERM;

  mutex_lock(&vhost_vapp_mutex);
  other = __vhost_vapp_get_dev(uuid);
  if (other && other != vapp) {
    mutex_unlock(&vhost_vapp_mutex);
    return -EADDRINUSE;
  }

  if (vapp->conf.uuid)
    hash_del_rcu(&vapp->node);

  vapp->conf.uuid = uuid;
  hash_add_rcu(vhost_vapp_hash, &vapp->node, vapp->conf.uuid);
  mutex_unlock(&vhost_vapp_mutex);

  pr_info("set uuid %llu\n", uuid);

  vapp->conf.uuid = uuid;
  return 0;
}

static long vhost_vapp_dev_set_multiqueue(struct vhost_vapp_dev* vapp, unsigned long arg)
{
  u64 multiqueues;
  struct vhost_vapp_dtq* ddtqs;
  struct vhost_virtqueue** vqs;
  size_t i;
  ssize_t ret;

  if (copy_from_user(&multiqueues, (void*)arg, sizeof(multiqueues)))
    return -EFAULT;

  pr_info("start set multiqueue to %llu\n", multiqueues);
  if (multiqueues == 0 || multiqueues > VIO_VAPP_NR_DATA_DTQS_MAX)
    return -EINVAL;

  ddtqs = vmalloc(multiqueues * sizeof(struct vhost_vapp_dtq));
  if (!ddtqs)
    return -ENOMEM;

  vqs = kmalloc_array(multiqueues * 2 + 2, sizeof(*vqs), GFP_KERNEL);
  if (!vqs) {
    ret = -ENOMEM;
    goto malloc_failed;
  }

  vqs[__VAPP_VQ_RX] = &vapp->cdtq.vqs[__VAPP_VQ_RX];
  vqs[__VAPP_VQ_TX] = &vapp->cdtq.vqs[__VAPP_VQ_TX];

  for (i = 0; i < multiqueues; i++) {
    struct vhost_vapp_dtq* dtq = &ddtqs[i];
    vhost_vapp_dtq_init_data(dtq);
    vhost_vapp_dtq_set_backend(dtq, vapp);
    vhost_vapp_dtq_set_queue(dtq, &vapp->dqueue);
    vqs[__VAPP_VQ_RX + (i + 1) * 2] = &dtq->vqs[__VAPP_VQ_RX];
    vqs[__VAPP_VQ_TX + (i + 1) * 2] = &dtq->vqs[__VAPP_VQ_TX];
  }

  for (i = 0; i < multiqueues * 2 + 2; i++)
    pr_info("register vhost queue: %p", vqs[i]);

  vhost_dev_init(&vapp->dev, vqs, multiqueues * 2 + 2, UIO_MAXIOV, VHOST_VAPP_PKT_WEIGHT, 0, true, NULL);

  vapp->conf.multiqueues = multiqueues;
  vapp->ddtqs = ddtqs;
  pr_info("set multiqueue to %llu\n", multiqueues);

  return 0;

 malloc_failed:
  kvfree(ddtqs);
  return ret;
}

static int vhost_vapp_dev_stop(struct vhost_vapp_dev* vapp)
{
  int ret;
  size_t i;

  mutex_lock(&vapp->dev.mutex);
  ret = vhost_dev_check_owner(&vapp->dev);
  if (ret)
    goto owner_fail;

  if (vapp->ddtqs)
    for (i = 0; i < vapp->conf.multiqueues; i++)
      vhost_vapp_dtq_stop(&vapp->ddtqs[i]);

 owner_fail:
  mutex_unlock(&vapp->dev.mutex);
  return ret;
}

static int vhost_vapp_dev_start(struct vhost_vapp_dev* vapp)
{
  int ret;
  size_t i;

  if (!vapp->ddtqs)
    return -EPERM;

  mutex_lock(&vapp->dev.mutex);
  ret = vhost_dev_check_owner(&vapp->dev);
  if (ret)
    goto owner_fail;

  ret = vhost_vapp_dtq_start(&vapp->cdtq);
  if (ret)
    goto vq_failed;
  for (i = 0; i < vapp->conf.multiqueues; i++) {
    struct vhost_vapp_dtq* dtq = &vapp->ddtqs[i];
    ret = vhost_vapp_dtq_start(dtq);
    if (ret)
      goto vq_failed;
  }
  pr_info("started queues\n");

  // TODO: work queue
  vhost_vapp_dtq_work_queue(&vapp->cdtq);
  for (i = 0; i < vapp->conf.multiqueues; i++) {
    struct vhost_vapp_dtq* dtq = &vapp->ddtqs[i];
    vhost_vapp_dtq_work_queue(dtq);
  }
  pr_info("worked queues\n");

  mutex_unlock(&vapp->dev.mutex);
  return 0;

 vq_failed:
  vhost_vapp_dtq_stop(&vapp->cdtq);
  for (i = 0; i < vapp->conf.multiqueues; i++) {
    struct vhost_vapp_dtq* dtq = &vapp->ddtqs[i];
    vhost_vapp_dtq_stop(dtq);
  }
 owner_fail:
  mutex_unlock(&vapp->dev.mutex);
  return ret;
}

static long vhost_vapp_dev_set_running(struct vhost_vapp_dev* vapp, unsigned long arg)
{
  u64 start;
  if (copy_from_user(&start, (void*)arg, sizeof(start)))
    return -EFAULT;

  pr_info("set start status %llu\n", start);

  if (start)
    return vhost_vapp_dev_start(vapp);
  else
    return vhost_vapp_dev_stop(vapp);
}

static long vhost_vapp_dev_ioctl(struct file* fp, unsigned int cmd, unsigned long arg)
{
  struct vhost_vapp_dev* vapp = fp->private_data;
  int ret = -ENOIOCTLCMD;
  switch (cmd) {
    case VHOST_VAPP_IO_SET_RUNNING:
      return vhost_vapp_dev_set_running(vapp, arg);
    case VHOST_VAPP_IO_SET_UUID:
      return vhost_vapp_dev_set_uuid(vapp, arg);
    case VHOST_VAPP_IO_SET_MULTIQUEUES:
      return vhost_vapp_dev_set_multiqueue(vapp, arg);
    case VIO_VAPP_IO_OPEN_REQ:
    case VIO_VAPP_IO_CLOSE_REQ:
      return -ENOIOCTLCMD;
    case VHOST_GET_FEATURES:
      return vhost_vapp_get_features(arg);
    case VHOST_SET_FEATURES:
      return vhost_vapp_set_features(vapp, arg);
    case VHOST_GET_BACKEND_FEATURES:
    case VHOST_SET_BACKEND_FEATURES:
      return -ENOIOCTLCMD;
    default:
      mutex_lock(&vapp->dev.mutex);
      ret = vhost_dev_ioctl(&vapp->dev, cmd, (void*)arg);
      if (ret == -ENOIOCTLCMD)
        ret = vhost_vring_ioctl(&vapp->dev, cmd, (void*)arg);
      else
        vhost_vapp_flush(vapp);
      mutex_unlock(&vapp->dev.mutex);
      return ret;
  }
  return 0;
}

static int vhost_vapp_dev_open(struct inode* inode, struct file* fp)
{
  struct vhost_vapp_dev* vapp = vhost_vapp_dev_alloc();
  if (!vapp)
    return -ENOMEM;
  fp->private_data = vapp;
  return 0;
}

static int vhost_vapp_dev_release(struct inode* inode, struct file* fp)
{
  struct vhost_vapp_dev* vapp = fp->private_data;

  mutex_lock(&vhost_vapp_mutex);
  if (vapp->conf.uuid)
    hash_del_rcu(&vapp->node);
  mutex_unlock(&vhost_vapp_mutex);

  synchronize_rcu();

  pr_info("removed dev %llu\n", vapp->conf.uuid);

  if (vapp->ddtqs) {
    vhost_vapp_dev_stop(vapp);
    vhost_vapp_flush(vapp);
    vhost_dev_stop(&vapp->dev);

    vhost_vapp_dev_destroy(vapp);

    vhost_dev_cleanup(&vapp->dev);
    kvfree(vapp->dev.vqs);
  }

  vhost_vapp_dev_free(vapp);

  return 0;
}

static const struct file_operations vhost_vapp_fops = {
  .owner          = THIS_MODULE,
  .open           = vhost_vapp_dev_open,
  .release        = vhost_vapp_dev_release,
  .llseek         = noop_llseek,
  .unlocked_ioctl = vhost_vapp_dev_ioctl,
  .compat_ioctl   = compat_ptr_ioctl,
};

static struct miscdevice vhost_vapp_misc = {
  .minor  = VHOST_VAPP_DEVICE_MINOR,
  .name   = VHOST_VAPP_DEVICE_NAME,
  .fops   = &vhost_vapp_fops,
  .mode   = VHOST_VAPP_DEVICE_MODE,
};


static int __init vhost_vapp_init(void)
{
  int ret = vio_dtq_init();
  if (ret < 0)
    return ret;

  ret = misc_register(&vhost_vapp_misc);
  if (ret)
    goto misc_failed;

  ret = vio_vapp_init();
  if (ret)
    goto vio_vapp_failed;

  return 0;

 vio_vapp_failed:
  misc_deregister(&vhost_vapp_misc);
 misc_failed:
  vio_dtq_exit();
  return ret;
}

static void __exit vhost_vapp_exit(void)
{
  vio_vapp_exit();
  misc_deregister(&vhost_vapp_misc);
}


module_init(vhost_vapp_init);
module_exit(vhost_vapp_exit);
MODULE_AUTHOR("Inochi Amaoto");
MODULE_DESCRIPTION("vhost virtual application backend driver");
MODULE_LICENSE("GPL v2");
MODULE_ALIAS("devname:vhost-vapp");
